package com.efeichong.common.config.flowLimit;

import com.efeichong.uuid.IdGenerator;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * @author lxk
 * @date 2022/8/22
 * @description 限流处理(滑动窗口)
 */
@Slf4j
@Component
public class SlideWindowFlowLimitHandler {

    private static final String LUA_SCRIPT = "-- 下标从 1 开始 获取key\n" +
            "local key = KEYS[1]\n" +
            "-- 下标从 1 开始 获取参数,  Lua  tonumber这个函数会尝试将它的参数转换为数字\n" +
            "local now = tonumber(ARGV[1]) \t\t\t   -- 当前时间戳\n" +
            "local slideExpiredTime = tonumber(ARGV[2]) -- 当前key 滑动窗口之外过期时间\n" +
            "local max = tonumber(ARGV[3])              -- 最大限流次数\n" +
            "local expiredTime = tonumber(ARGV[4])      -- 缓存过期时间\n" +
            "local random = ARGV[5]      -- 缓存过期时间\n" +
            "\n" +
            "-- 移除滑动窗口之外的数据 Redis ZREMRANGEBYSCORE 命令移除有序集key中，所有score值介于min和max之间(包括等于min或max)的成员。\n" +
            "-- 移除指定分数区间内的所有元素，expired 即已经过期的 score, 根据当前时间毫秒数 - 超时毫秒数，得到过期时间 slideExpiredTime\n" +
            "redis.call('zremrangebyscore', key, 0, slideExpiredTime)\n" +
            "\n" +
            "-- 每次访问均重新设置 zset 的过期时间，单位毫秒\n" +
            "redis.call(\"pexpire\", key, expiredTime)\n" +
            "\n" +
            "-- 获取 zset 中的当前元素个数 , Redis ZCARD 命令用于返回有序集的成员个数\n" +
            "local current = tonumber(redis.call('zcard', key))\n" +
            "local next = current + 1\n" +
            "-- 判断当前集合内个数+1 后是否超出最大限制次数\n" +
            "if next > max then\n" +
            "  -- 达到限流大小 返回 0\n" +
            "  return 0\n" +
            "-- 否则添加进入zset集合\n" +
            "else\n" +
            "  -- 往 zset 中添加一个 Score、value均为当前时间戳的元素，[score,value]\n" +
            "  redis.call(\"zadd\", key, random, random)\n" +
            "\n" +
            "  -- PEXPIRE 命令和 EXPIRE 命令的作用类似，但是它以毫秒为单位设置 key 的生存时间，而不像 EXPIRE 命令那样，以秒为单位\n" +
            "  -- 返回当前次数\n" +
            "  return next\n" +
            "end";

    @Autowired
    private StringRedisTemplate redisTemplate;

    public boolean flow(String key, int capacity) {
        // 假设以下是通过注解获取到的配置, 或者其他地方获取到的配置
        // (限流时间内允许的) 最大请求
        long max = capacity;
        // 滑动窗口时间(限流时间范围)
        long period = 1L;
        // 限流时间单位, 默认秒
        TimeUnit timeUnit = TimeUnit.SECONDS;
        // 缓存的过期时间, 转成毫秒
        long expiredTime = timeUnit.toMillis(period);
        // 当前时间戳
        long currentTimeMillis = System.currentTimeMillis();
        // 当前时间 - 间隔时间 = 滑动窗口外数据过期时间
        long slideExpiredTime = currentTimeMillis - expiredTime;
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript(LUA_SCRIPT, Long.class);
        // 执行lua脚本判断是否达到限流    脚本对象,   keys集合,  参数集合(根据脚本索引放)
        String random = IdGenerator.getIdStr();
        Long result = redisTemplate.execute(redisScript, Collections.singletonList(key),
                String.valueOf(currentTimeMillis), String.valueOf(slideExpiredTime), String.valueOf(max),
                String.valueOf(expiredTime),random);
        result = Optional.ofNullable(result).orElse(-1L);
        log.info("result={}", result);
        return result != 0;
    }
}
